"""
# Licensed under the MIT license.

蒙特卡洛树搜索(MCTS)算法核心实现

本模块提供了MCTS算法的基础框架实现，包括：
1. 抽象节点类定义 - 提供统一接口规范
2. 搜索器类实现 - 执行完整的MCTS搜索过程
3. UCT公式实现 - 平衡探索与利用

主要功能：
- 提供通用的MCTS算法框架，可被特定领域继承扩展
- 实现四个核心阶段：Selection、Expansion、Simulation、Backpropagation
- 支持多种探索权重调度策略

原始代码来源: Luke Harold Miles, July 2019
修改扩展: rStar团队
"""

from abc import ABC, abstractmethod  # 用于创建抽象基类和抽象方法
from collections import defaultdict  # 提供默认值的字典，避免键不存在的异常
from typing import Dict, List  # 类型注解，提高代码可读性和类型安全
import math  # 数学函数库，用于UCT公式中的对数和平方根计算
import random  # 随机数生成，用于模拟阶段的随机选择


# 全局节点计数器，用于为每个节点分配唯一ID
node_cnt = 0


def verbose_print(s: str, verbose: bool):
    """
    条件打印函数，仅在verbose模式下输出信息
    
    Args:
        s (str): 待打印的字符串
        verbose (bool): 是否处于详细输出模式
    """
    if verbose:
        print(s)


class MCTS_Node(ABC):
    """
    MCTS节点抽象基类
    
    定义了MCTS树中每个节点必须实现的接口，包括：
    - 子节点生成 (find_children)
    - 终止状态判断 (is_terminal)
    - 奖励计算 (calculate_reward)
    - 反向传播控制 (skip_backprop)
    
    派生类必须实现这些抽象方法以适应特定领域问题
    """

    def __init__(self) -> None:
        """
        初始化节点，分配唯一ID
        """
        super().__init__()

        global node_cnt
        self.id = node_cnt  # 节点唯一标识符
        node_cnt += 1

        self.rollout_id = None  # 记录节点所属的rollout轮次

    def set_rollout_id(self, rollout_id: int):
        """
        设置节点所属的rollout轮次ID
        
        Args:
            rollout_id (int): rollout轮次标识符
        """
        self.rollout_id = rollout_id

    @abstractmethod
    def find_children(self, rollout_id: int):
        """
        生成当前节点的所有可能子节点
        
        Args:
            rollout_id (int): 当前rollout轮次ID
            
        Returns:
            List[MCTS_Node]: 子节点列表
        """
        raise NotImplementedError

    @abstractmethod
    def is_terminal(self):
        """
        判断当前节点是否为终止节点
        
        Returns:
            bool: 如果是终止节点则返回True，否则返回False
        """
        raise NotImplementedError

    @abstractmethod
    def calculate_reward(self):
        """
        计算终止节点的奖励值
        
        通常对于终止节点，计算其奖励值用于反向传播
        例如：1=胜利，0=失败，0.5=平局
        
        Returns:
            float: 节点的奖励值
        """
        raise NotImplementedError

    @abstractmethod
    def skip_backprop(self):
        """
        判断是否跳过该节点的反向传播
        
        某些节点(如根节点)可能不需要参与奖励的累积计算
        
        Returns:
            bool: 如果为True则在反向传播时跳过该节点
        """
        raise NotImplementedError


class MCTS_Searcher:
    """
    蒙特卡洛树搜索器实现
    
    执行完整的MCTS搜索过程，维护节点统计信息，
    并通过UCT公式平衡探索与利用
    """

    def __init__(
        self,
        exploration_weight: float,
        weight_scheduler: str,
        num_rollouts: int,
        discount: float,
        verbose: bool = False,
    ):
        """
        初始化MCTS搜索器
        
        Args:
            exploration_weight (float): UCT公式中的探索权重
            weight_scheduler (str): 权重调度策略 ('exp'/'lin'/'const')
            num_rollouts (int): 总rollout轮数
            discount (float): 奖励折扣因子
            verbose (bool): 是否输出详细日志
        """
        # 节点统计数据结构
        self.Q: Dict[MCTS_Node, float] = defaultdict(lambda: 0.0)  # 节点累积奖励值
        self.N: Dict[MCTS_Node, int] = defaultdict(lambda: 0)  # 节点访问次数
        self.parent2children: Dict[MCTS_Node, List[MCTS_Node]] = dict()  # 父节点到子节点的映射

        # 已探索节点集合(已扩展+已模拟)，具有Q和N值的节点
        self.explored_nodes = set()

        # 算法参数
        self.exploration_weight = exploration_weight  # UCT探索权重
        self.weight_scheduler = weight_scheduler      # 权重调度策略
        self.num_rollouts = num_rollouts              # 总rollout轮数
        self.discount = discount                      # 奖励折扣因子

        self.verbose = verbose  # 详细输出控制

        # 重置全局节点计数器
        global node_cnt
        node_cnt = 0

    def do_rollout(self, root_node: MCTS_Node, rollout_id: int):
        """
        执行一轮完整的MCTS rollout
        
        包含四个核心步骤：
        1. Selection: 选择最有潜力的节点
        2. Expansion: 扩展选中节点
        3. Simulation: 从扩展节点模拟到终端状态
        4. Backpropagation: 反向传播奖励值
        
        Args:
            root_node (MCTS_Node): 搜索树的根节点
            rollout_id (int): 当前rollout轮次ID
            
        Returns:
            MCTS_Node: rollout结束的节点，通常是终端节点
        """
        verbose_print("==> Selecting a node...", self.verbose)
        path_1 = self._select(root_node, rollout_id)
        leaf = path_1[-1]
        verbose_print(f"==> Expanding node {leaf.id}...", self.verbose)
        self._expand(leaf, rollout_id)
        verbose_print(f"==> Simulating node {leaf.id}...", self.verbose)
        path_2 = self._simulate(leaf, rollout_id)
        verbose_print(f"==> Backpropagating...", self.verbose)
        self._backpropagate(path_1 + path_2)
        try:
            return path_2[-1]  # 返回模拟路径的最后一个节点
        except:
            return path_1[-1]  # 如果模拟路径为空，返回选择路径的最后一个节点

    def _select(self, node: MCTS_Node, rollout_id: int) -> List[MCTS_Node]:
        """
        选择阶段：找到最有潜力的未探索节点
        
        从根节点开始，根据以下规则递归选择：
        1. 如果节点没有子节点，直接选择该节点
        2. 如果节点有未探索的子节点，随机选择一个未探索子节点
        3. 如果所有子节点都已探索，使用UCT公式选择最佳子节点
        
        Args:
            node (MCTS_Node): 当前节点
            rollout_id (int): 当前rollout轮次ID
            
        Returns:
            List[MCTS_Node]: 从根节点到选中节点的路径
        """
        path = []
        while True:
            path.append(node)
            # 情况1：节点没有子节点，选择节点本身
            if node not in self.parent2children.keys():
                return path

            # 情况2：节点有未探索的子节点，随机选择一个
            unexplored = [n for n in self.parent2children[node] if n not in self.explored_nodes]
            if unexplored:
                n = random.choice(unexplored)
                path.append(n)
                return path

            # 情况3：所有子节点都已探索，使用UCT选择最佳子节点
            node = self._uct_select(node, rollout_id)

    def _expand(self, node: MCTS_Node, rollout_id: int):
        """
        扩展阶段：生成选中节点的子节点
        
        如果节点已经扩展过或是终端节点，则不进行扩展
        否则调用节点的find_children方法生成子节点
        
        Args:
            node (MCTS_Node): 待扩展的节点
            rollout_id (int): 当前rollout轮次ID
        """
        if node in self.explored_nodes:
            return  # 节点已扩展，无需再次扩展

        if node.is_terminal():
            self.explored_nodes.add(node)
            return  # 终端节点不可扩展

        self.parent2children[node] = node.find_children(rollout_id)

    def _simulate(self, node: MCTS_Node, rollout_id: int) -> List[MCTS_Node]:
        """
        模拟阶段：从扩展节点随机模拟到终端状态
        
        从给定节点开始，随机选择子节点直到达到终端状态
        
        Args:
            node (MCTS_Node): 模拟起始节点
            rollout_id (int): 当前rollout轮次ID
            
        Returns:
            List[MCTS_Node]: 模拟过程中访问的节点路径
        """
        path = []
        cur_node = node
        while True:
            if cur_node.is_terminal():
                self.explored_nodes.add(node)
                return path

            if cur_node not in self.parent2children.keys():
                self.parent2children[cur_node] = cur_node.find_children(rollout_id)

            cur_node = random.choice(self.parent2children[cur_node])  # 随机选择一个子节点
            path.append(cur_node)

    def _backpropagate(self, path: List[MCTS_Node]):
        """
        反向传播阶段：将奖励值传播回路径上的所有节点
        
        计算终端节点的奖励，并更新路径上所有节点的统计信息
        
        Args:
            path (List[MCTS_Node]): 从根节点到终端节点的完整路径
        """
        leaf = path[-1]
        reward = leaf.calculate_reward()
        for node in reversed(path):
            self.Q[node] += reward
            self.N[node] += 1
            self.explored_nodes.add(node)

    def _get_weight(self, rollout_id: int):
        """
        根据调度策略计算当前rollout的探索权重
        
        支持三种权重调度策略：
        - exp: 指数衰减 (从exploration_weight衰减到0.1*exploration_weight)
        - lin: 线性衰减 (从exploration_weight衰减到0.1*exploration_weight)
        - const: 恒定权重 (保持exploration_weight不变)
        
        Args:
            rollout_id (int): 当前rollout轮次ID
            
        Returns:
            float: 当前轮次的探索权重
        """
        # 权重调度：从初始探索权重逐渐降低到0.1倍初始权重
        if self.weight_scheduler == "exp":
            return self.exploration_weight * (0.1 ** (rollout_id / self.num_rollouts))
        elif self.weight_scheduler == "lin":
            return self.exploration_weight * (1 - 0.9 * (rollout_id / self.num_rollouts))
        elif self.weight_scheduler == "const":
            return self.exploration_weight
        else:
            return self.exploration_weight  # 默认返回恒定权重

    def _uct_select(self, node: MCTS_Node, rollout_id: int):
        """
        使用UCT公式选择最佳子节点
        
        在所有已探索的子节点中，选择UCT值最大的节点
        
        Args:
            node (MCTS_Node): 父节点
            rollout_id (int): 当前rollout轮次ID
            
        Returns:
            MCTS_Node: UCT值最大的子节点
        """
        # 确保所有子节点都已被探索过
        assert all(n in self.explored_nodes for n in self.parent2children[node])

        return max(
            self.parent2children[node], key=lambda n: self._compute_uct(parent_node=node, node=n, rollout_id=rollout_id)
        )

    def _compute_uct(self, parent_node: MCTS_Node, node: MCTS_Node, rollout_id: int):
        """
        计算节点的UCT值
        
        UCT公式 = 平均奖励 + 探索权重 * sqrt(ln(父节点访问次数)/节点访问次数)
        
        Args:
            parent_node (MCTS_Node): 父节点
            node (MCTS_Node): 当前节点
            rollout_id (int): 当前rollout轮次ID
            
        Returns:
            float: 节点的UCT值，用于选择最佳子节点
        """
        if parent_node is None:  # 无效UCT：节点是根节点
            return 666
        else:
            if self.N[node] == 0:  # 无效UCT：节点未被探索过
                return 999
            else:
                weight = self._get_weight(rollout_id)
                # UCT公式：平均奖励 + 探索项
                if weight is None:  # 修复None不支持*运算符的错误
                    weight = self.exploration_weight
                return self.Q[node] / self.N[node] + weight * math.sqrt(math.log(self.N[parent_node]) / self.N[node])
